{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9aa64542",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T07:19:09.652991Z",
     "start_time": "2022-11-10T07:19:09.647856Z"
    }
   },
   "source": [
    "## extract data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16a6806d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T07:50:32.084849Z",
     "start_time": "2022-11-10T07:50:32.079142Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba7ad94b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-09T11:36:43.120189Z",
     "start_time": "2022-11-09T11:36:43.116217Z"
    }
   },
   "outputs": [],
   "source": [
    "json_file = 'dataset/tool/category_info_gen_all.json'\n",
    "tar_path = '/data1/ImageNet/trainset'\n",
    "extract_path = '/data1/style_transfer/rec'\n",
    "\n",
    "folder_class = {\n",
    "    '人物': '000',\n",
    "    '动物': '001',\n",
    "    '植物': '002',\n",
    "    '自然风景': '003',\n",
    "    '人造物': '004',\n",
    "    '建筑': '005'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "997b41bd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-09T11:36:43.126599Z",
     "start_time": "2022-11-09T11:36:43.123128Z"
    }
   },
   "outputs": [],
   "source": [
    "def mkdir(path):\n",
    "    folder = os.path.exists(path)\n",
    "    if not folder:\n",
    "        os.makedirs(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "164bd660",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-09T11:36:43.135249Z",
     "start_time": "2022-11-09T11:36:43.128850Z"
    }
   },
   "outputs": [],
   "source": [
    "info_all = json.loads(open(json_file, 'r').read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "261a3daa",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2022-11-09T11:36:00.156Z"
    }
   },
   "outputs": [],
   "source": [
    "for item in info_all:\n",
    "    desc = item['desc']\n",
    "    item_class = item['content']\n",
    "    if desc not in folder_class:\n",
    "        continue\n",
    "    folder_id = folder_class[desc]\n",
    "    outpath = os.path.join(extract_path, folder_id)\n",
    "    if os.path.exists(outpath):\n",
    "        continue\n",
    "    mkdir(outpath)\n",
    "    print(f'extracting {desc} to {outpath}')\n",
    "    pbar = tqdm(item_class)\n",
    "    for ii in pbar:\n",
    "        pbar.set_description('image label %s - %s' %\n",
    "                             (ii['cn_string'], ii['imagenet_string']))\n",
    "        tar_ii = os.path.join(tar_path, ii['imagenet_label'] + '.tar')\n",
    "        ext_cmd = f'tar -xf {tar_ii} --directory={outpath}'\n",
    "        os.system(ext_cmd)\n",
    "#         print(ext_cmd, os.path.exists(tar_ii))\n",
    "#         time.sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1af6a9a0",
   "metadata": {},
   "source": [
    "## set split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b178bc28",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T07:46:59.445469Z",
     "start_time": "2022-11-10T07:46:59.437808Z"
    }
   },
   "outputs": [],
   "source": [
    "class_num = {\n",
    "    '000': 5000,\n",
    "    '001': 9000,\n",
    "    '002': 9000,\n",
    "    '003': 7000,\n",
    "    '004': 10000,\n",
    "    '005': 9000,\n",
    "}\n",
    "output_path = '/data1/style_transfer/rec_split'\n",
    "trainset_path = os.path.join(output_path, 'train')\n",
    "testset_path = os.path.join(output_path, 'test')\n",
    "src_path = '/data1/style_transfer/rec'\n",
    "testset_ratio = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb973cf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T08:03:10.908480Z",
     "start_time": "2022-11-10T08:03:10.900844Z"
    }
   },
   "outputs": [],
   "source": [
    "def copy_files(filelist, src, dst):\n",
    "    for item in tqdm(filelist):\n",
    "        if not item.lower().endswith('jpeg'):\n",
    "            continue\n",
    "        \n",
    "        src_f = os.path.join(src, item)\n",
    "        dst_f = os.path.join(dst, item)\n",
    "        cmd = 'cp %s %s'%(src_f, dst_f)\n",
    "        os.system(cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f268f666",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T08:18:12.307710Z",
     "start_time": "2022-11-10T08:07:48.842429Z"
    }
   },
   "outputs": [],
   "source": [
    "for name, num in class_num.items():\n",
    "    item_path_src = os.path.join(src_path, name)\n",
    "    item_path_dst_train = os.path.join(trainset_path, name)\n",
    "    item_path_dst_test = os.path.join(testset_path, name)\n",
    "    total = os.listdir(item_path_src)\n",
    "    random.shuffle(total)\n",
    "    total = total[:num]\n",
    "    split_num = int(num * testset_ratio)\n",
    "    test_set = total[:split_num]\n",
    "    train_set = total[split_num:]\n",
    "    print(f'split {name} train_set:{len(train_set)} test_set:{len(test_set)}')\n",
    "    mkdir(item_path_dst_train)\n",
    "    mkdir(item_path_dst_test)\n",
    "    print('copying train set')\n",
    "    copy_files(train_set, item_path_src, item_path_dst_train)\n",
    "    print('copying test set')\n",
    "    copy_files(test_set, item_path_src, item_path_dst_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dffd820d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
